from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from pydantic import Field,BaseModel

# 将文本分配到标签中
# 是指使用如下类标记文档： 情绪、语言、风格（正式、非正式等）、涵盖的主题、政治趋势

# 定义提示模板
tagging_prompt = ChatPromptTemplate.from_template(
    """
Extract the desired information from the following passage.

Only extract the properties mentioned in the 'Classification' function.

Passage:
{input}
"""
)

class Classification(BaseModel):
    sentiment: str = Field(..., enum=["happy", "neutral", "sad"]) # 情绪
    aggressiveness: int = Field(
        ...,
        description="describes how aggressive the statement is, the higher the number the more aggressive",
        enum=[1, 2, 3, 4, 5],
    ) # 风格
    language: str = Field(
        ..., enum=["spanish", "english", "french", "german", "italian"]
    ) # 语言

# 定义模型 with_structured_output方法，做为结构化输出
llm = ChatOpenAI(temperature=0, model="gpt-4o-mini").with_structured_output(
    Classification
)
chain = tagging_prompt | llm

result = chain.invoke({"input":"nice to meet you!"})

print(result)

